[WIP]: [Fix] Per-rank SHM + streaming 1× host peak + FX device placement - #64

…tion

Two tests using torch.multiprocessing.spawn with gloo backend:

1. test_shared_memory_overwrites_ep_shards:
   Reproduces the bug — local_rank=0 writes expert weights to a shared
   file, all other ranks read it, silently overwriting their own expert
   shards with rank 0's data.

2. test_ep_fix_preserves_per_rank_shards:
   Verifies the fix — when EP_SIZE > 1, the shared-memory path is
   skipped and each rank retains its own expert weights.

ENGINE_CONFIG__EP_SIZE may not be set if the host framework
(e.g. disagg_compute_runner) only sets EP_SIZE or configures
ep_size programmatically. Fall back to EP_SIZE env var before
defaulting to 1.

When EP_SIZE > 1, each rank holds a unique expert shard. The previous
fix skipped shared memory entirely and used pin_memory, which was
extremely slow for large models (~46GB per rank).

Now each rank writes its own shared-memory file to /dev/shm and
mmap-reads it back, preserving per-rank expert weights while keeping
the speed benefit of shared memory + pin_memory_in_place on
already-resident pages.

For EP_SIZE <= 1, the original rank-0-writes-all-read scheme is
retained (all ranks have identical weights).

Also updates the regression test to verify the per-rank shm path.

1. _force_cpu: skip GPU roundtrip for CPU tensors when fn only changes
   device (not dtype). Reduces peak host memory during model.cuda() by
   avoiding temporary CUDA host allocations for every parameter.

2. MAGI_OFFLOAD_SKIP_SHM: when set to "1", skip shared memory creation
   and pin_memory_in_place entirely. Params remain as regular CPU tensors.
   This allows OffloadExecutor to work on memory-constrained nodes (e.g.
   5090 with 512Gi container limit for 8x EP ranks) where the shm+pin
   overhead causes OOM.

When converting CPU example_value metadata to CUDA for Inductor, the
.to(device) call strips torch.nn.Parameter wrapping. This caused
OffloadExecutor to misidentify all model weights as regular input
tensors, loading all 43.6GB onto GPU simultaneously instead of
offloading per-submodule — OOM on 5090 (31GB VRAM).

Re-wrap the converted FakeTensor in nn.Parameter to preserve type info.

- OffloadExecutor: log per-step H2D vs compute breakdown when MAGI_OFFLOAD_DEBUG=1
  (cuda.synchronize between prefetch and compute for accurate wall-clock split)
- _patch_cpu_offload_apply: support MAGI_OFFLOAD_PIN_BUDGET_GB env var
  Pin up to N GB of weights per rank via cudaHostRegister (no SHM copy)
  for faster async H2D while staying within host memory budget

Log start/end time for each rank during staggered pin, plus OS memlock
limit. Helps debug slow NFS page faults during cudaHostRegister.

Replace sequential 1-rank-at-a-time pinning with parallel waves.
Auto-detects max concurrent ranks: total_ram/2 / per_rank_param_size.
Override via MAGI_OFFLOAD_PIN_CONCURRENCY env var.

512GB node, 43.36GB/rank → concurrency=5, 2 waves instead of 8.
Expected pin time: ~6min vs ~22min sequential.

…-corruption

The per-submodule cuda.synchronize() barriers prevented H2D/compute
pipeline overlap, reducing production throughput. Profiling data has
been collected; this debug scaffolding is no longer needed.

- Revert offload_warpper.py (all changes were unused imports after debug removal)
- Remove dead offload() __dict__ branch (call sites only pass tuple/dict)
- Extract 80-line inline pin logic into _staggered_pin_memory() helper
- magi_backend.py: use module-level os/magi_logger instead of inline imports
- Use %-formatting instead of f-string for logger calls

…fload_apply

- _shm_path(): centralize /dev/shm path construction
- _pack_params_flat(): copy named tensors into contiguous buffer
- _split_flat_to_params(): split flat buffer back to named param views
- _create_shm_tensor(): create mmap file + pack in one call
- Unify EP>1 and EP<=1 serialization (both use mmap now, remove numpy bf16/fp8 workaround)
- _patch_cpu_offload_apply SHM logic: ~135 lines -> ~45 lines

- Merge per_rank_shm and shared branches into single helper
- Collapse 3-way dispatch to 2-way (skip_shm vs materialize)
- Single cleanup point for del/gc.collect()
- Eliminate duplicated pin/append/split/remove/load_state_dict

…s.environ

…-corruption

…mpat)

Extract get_cpu_gloo_group() into utils/dist_utils.py, shared by
_all_ranks_same_weights and _get_cost_sync_group. Fixes RuntimeError
on NCCL-default process groups where CPU tensors are rejected.

… to per_rank=True)